/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.runners.gaia;



import com.bff.gaia.unified.sdk.Pipeline;

import com.bff.gaia.unified.sdk.Pipeline.PipelineExecutionException;

import com.bff.gaia.unified.sdk.PipelineResult;

import com.bff.gaia.unified.sdk.PipelineRunner;

import com.bff.gaia.unified.sdk.options.PipelineOptions;

import com.bff.gaia.unified.sdk.options.PipelineOptionsFactory;

import com.bff.gaia.unified.sdk.util.UserCodeException;



/** Test Gaia runner. */

public class TestGaiaRunner extends PipelineRunner<PipelineResult> {



  private GaiaRunner delegate;



  private TestGaiaRunner(GaiaPipelineOptions options) {

    options.setRunner(TestGaiaRunner.class);

    options.setShutdownSourcesOnFinalWatermark(true);

    if (options.getParallelism() == -1) {

      // Limit parallelism to avoid too much memory consumption during local execution

      options.setParallelism(1);

    }

    this.delegate = GaiaRunner.fromOptions(options);

  }



  public static TestGaiaRunner fromOptions(PipelineOptions options) {

    GaiaPipelineOptions gaiaOptions = options.as(GaiaPipelineOptions.class);

    return new TestGaiaRunner(gaiaOptions);

  }



  public static TestGaiaRunner create(boolean streaming) {

    GaiaPipelineOptions gaiaOptions = PipelineOptionsFactory.as(GaiaPipelineOptions.class);

    gaiaOptions.setStreaming(streaming);

    return fromOptions(gaiaOptions);

  }



  @Override

  public PipelineResult run(Pipeline pipeline) {

    try {

      return delegate.run(pipeline);

    } catch (Throwable t) {

      // Special case hack to pull out assertion errors from PAssert; instead there should

      // probably be a better story along the lines of UserCodeException.

      UserCodeException innermostUserCodeException = null;

      Throwable current = t;

      for (; current.getCause() != null; current = current.getCause()) {

        if (current instanceof UserCodeException) {

          innermostUserCodeException = (UserCodeException) current;

        }

      }

      if (innermostUserCodeException != null) {

        current = innermostUserCodeException.getCause();

      }

      if (current instanceof AssertionError) {

        throw (AssertionError) current;

      }

      throw new PipelineExecutionException(current);

    }

  }



  public PipelineOptions getPipelineOptions() {

    return delegate.getPipelineOptions();

  }

}